package huu.phong.sshchecker;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.Scanner;
import java.util.StringTokenizer;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import com.jcraft.jsch.JSch;

public class SSHChecker {
	public static void main(String[] args) throws IOException {
		if (args.length < 2) {
			System.out.println("Enter some args");
			System.exit(0);
		}
		
		final Properties properties = new Properties();
		FileInputStream conf = new FileInputStream("sshchecker.conf");
		properties.load(conf);
		conf.close();
		
		final List<String> countries = new ArrayList<String>();
		StringTokenizer countryToken = new StringTokenizer(properties.getProperty("countries"), ",");
		while (countryToken.hasMoreTokens()) {
			String country = countryToken.nextToken();
			if (!countries.contains(country)) countries.add(country.toLowerCase().trim());
		}
		
		int thread = Integer.parseInt(args[0]);
		final int timeout = Integer.parseInt(args[1]);
		String filename = "ssh.txt";
		if (args.length >= 3) {
			filename = args[2];
		}
		final boolean isGetCountry = args.length == 4;
		
		InputStream in = new FileInputStream(filename);
		
		String hosts = streamToString(in);
		
		List<SSHHost> sshHosts = new ArrayList<SSHHost>();
		
		StringTokenizer token = new StringTokenizer(hosts, "\r\n");
		while (token.hasMoreTokens()) {
			StringTokenizer hostToken = new StringTokenizer(token.nextToken(), "|;/,");
			if (hostToken.countTokens() < 3) continue;
			SSHHost host = new SSHHost();
			host.host = hostToken.nextToken().trim();
			host.username = hostToken.nextToken().trim();
			host.password = hostToken.nextToken().trim();
			if (sshHosts.contains(host)) continue;
			sshHosts.add(host);
		}
		final JSch jsch = new JSch();
		final SSHOutputStream out = new SSHOutputStream(filename);
		final IPToLocation ipToLocation = new IPToLocation();
		final LongCount count = new LongCount(sshHosts.size());
		ExecutorService executorService = Executors.newFixedThreadPool(thread);
		for (final SSHHost host : sshHosts) {
			executorService.submit(new Runnable() {
				@Override
				public void run() {
					Checker checker = new Checker(host.username, host.password, host.host);
					if (checker.check(jsch, timeout)) {
						if (isGetCountry) {
							try {
								String country = ipToLocation.getRegion(host.host).trim();
								if (countries.isEmpty() || countries.contains(country.toLowerCase())) {
									out.println(host.toString() + "|" + country);
								}
								System.out.println(String.format("%s %4s %-40s --> %30s", count.plusOne().toString(), "Done", host.toString(), country));
							} catch (IOException e) {
								//out.println(host.toString() + "|Unknow");
								//System.out.println("Done " + host.toString() + " --> Unknow");
							}
						} else {
							out.println(host.toString());
							System.out.println(count.plusOne().toString() + " Done " + host.toString());
						}
					} else {
						System.out.println(String.format("%s %-4s %-40s", count.plusOne().toString(), "Die", host.toString()));
					}
				}
			});
		}
		executorService.shutdown();
		if (executorService.isTerminated()) {
			System.out.println("DONE");
		}
	}
	
	public static String streamToString(InputStream in) throws IOException{
		Scanner scanner = new Scanner(in, "UTF-8").useDelimiter("\\A");
		String ret = scanner.hasNext() ? scanner.next(): "";
		in.close();
		return ret;
	}
}
